{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keyword Spotting (KWS) using PyTorch Lightning \n",
    "\n",
    "We illustrate how to use PyTorch Lightning to train a model for keyword spotting. We also show how to build a model that is data agnostic, i.e. it can be trained on any dataset. A separate module handles the datasets and dataloaders.\n",
    "\n",
    "Using Experience, Task and Performance concept, KWS can be described as follows:\n",
    "\n",
    "1. **Experience:** The model is trained on a dataset of audio files containing 1sec speech samples, each is a single keyword. There are 35 distinct words in the dataset such as `yes`, `no`, `right`, `left` and so forth. However, 2 additional categories are added to the dataset: `silence` and `unknown`. `silence` is when there is silence or non-word audio audio activity such as background noise. `unknown` is when there is a keyword but not one of those 35 distinct words.\n",
    "\n",
    "2. **Task:** The model is trained to classify keywords in audio samples. We use a modified ResNet18 model.\n",
    "\n",
    "3. **Performance:** We measure the performance of the model by calculating the accuracy of the model on the test set. \n",
    "\n",
    "For simplicity, we do not use the validation set for hyperparameter tuning.\n",
    "\n",
    "We use KWS Version 2. The original KWS paper can be found [here](https://arxiv.org/pdf/1804.03209.pdf).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pytorch-lightning in /home/rowel/anaconda3/lib/python3.7/site-packages (1.6.0)\n",
      "Requirement already satisfied: torch>=1.8.* in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (1.11.0)\n",
      "Requirement already satisfied: packaging>=17.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (21.3)\n",
      "Requirement already satisfied: typing-extensions>=4.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (4.1.1)\n",
      "Requirement already satisfied: numpy>=1.17.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (1.21.5)\n",
      "Requirement already satisfied: tqdm>=4.41.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (4.62.3)\n",
      "Requirement already satisfied: PyYAML>=5.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (6.0)\n",
      "Requirement already satisfied: pyDeprecate<0.4.0,>=0.3.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (0.3.2)\n",
      "Requirement already satisfied: tensorboard>=2.2.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (2.4.1)\n",
      "Requirement already satisfied: torchmetrics>=0.4.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (0.7.3)\n",
      "Requirement already satisfied: fsspec[http]!=2021.06.0,>=2021.05.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (2022.1.0)\n",
      "Requirement already satisfied: requests in /home/rowel/anaconda3/lib/python3.7/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.27.1)\n",
      "Requirement already satisfied: aiohttp in /home/rowel/anaconda3/lib/python3.7/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.7.4.post0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from packaging>=17.0->pytorch-lightning) (3.0.4)\n",
      "Requirement already satisfied: setuptools>=41.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (58.0.4)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.1.1)\n",
      "Requirement already satisfied: six>=1.10.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.16.0)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.6.0.post3)\n",
      "Requirement already satisfied: werkzeug>=0.11.15 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (2.0.3)\n",
      "Requirement already satisfied: grpcio>=1.24.3 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.32.0)\n",
      "Requirement already satisfied: protobuf>=3.6.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.19.4)\n",
      "Requirement already satisfied: google-auth<2,>=1.6.3 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.15.0)\n",
      "Requirement already satisfied: wheel>=0.26 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.37.1)\n",
      "Requirement already satisfied: absl-py>=0.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.12.0)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.4.1)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.1.0)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.2.8)\n",
      "Requirement already satisfied: rsa<4.1,>=3.1.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.0)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (1.3.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.4.8)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.3)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2021.10.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.26.8)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (3.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (5.1.0)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.6.3)\n",
      "Requirement already satisfied: async-timeout<4.0,>=3.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.0.1)\n",
      "Requirement already satisfied: chardet<5.0,>=2.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (4.0.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (21.4.0)\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Requirement already satisfied: torchmetrics in /home/rowel/anaconda3/lib/python3.7/site-packages (0.7.3)\n",
      "Requirement already satisfied: numpy>=1.17.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (1.21.5)\n",
      "Requirement already satisfied: torch>=1.3.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (1.11.0)\n",
      "Requirement already satisfied: pyDeprecate==0.3.* in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (0.3.2)\n",
      "Requirement already satisfied: packaging in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (21.3)\n",
      "Requirement already satisfied: typing-extensions in /home/rowel/anaconda3/lib/python3.7/site-packages (from torch>=1.3.1->torchmetrics) (4.1.1)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from packaging->torchmetrics) (3.0.4)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install pytorch-lightning --upgrade\n",
    "%pip install torchmetrics --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchaudio, torchvision\n",
    "import os\n",
    "import matplotlib.pyplot as plt \n",
    "import librosa\n",
    "import argparse\n",
    "import numpy as np\n",
    "import wandb\n",
    "from pytorch_lightning import LightningModule, Trainer, LightningDataModule, Callback\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from torchmetrics.functional import accuracy\n",
    "from torchvision.transforms import ToTensor\n",
    "from torchaudio.datasets import SPEECHCOMMANDS\n",
    "from torchaudio.datasets.speechcommands import load_speechcommands_item\n",
    "\n",
    "\n",
    "# use second GPU\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Silence and Unknown Datasets\n",
    "\n",
    "We create a custom `silence` dataset. The dataset randomly samples background audio supplied in the KWS dataset. These files are under the `_background_noise_` folder.\n",
    "\n",
    "We also create `unknown` dataset uses random audio samples from the train set but labelled as `unknown`. \n",
    "\n",
    "The creation of these 2 datasets is described in the [KWS paper](https://arxiv.org/pdf/1804.03209.pdf) and implemented below.\n",
    "\n",
    "We limit the number of samples to about the size of train dataset divided by 35 (the number of distinct words in the KWS dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SilenceDataset(SPEECHCOMMANDS):\n",
    "    def __init__(self, root):\n",
    "        super(SilenceDataset, self).__init__(root, subset='training')\n",
    "        self.len = len(self._walker) // 35\n",
    "        path = os.path.join(self._path, torchaudio.datasets.speechcommands.EXCEPT_FOLDER)\n",
    "        self.paths = [os.path.join(path, p) for p in os.listdir(path) if p.endswith('.wav')]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        index = np.random.randint(0, len(self.paths))\n",
    "        filepath = self.paths[index]\n",
    "        waveform, sample_rate = torchaudio.load(filepath)\n",
    "        return waveform, sample_rate, \"silence\", 0, 0\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.len\n",
    "\n",
    "class UnknownDataset(SPEECHCOMMANDS):\n",
    "    def __init__(self, root):\n",
    "        super(UnknownDataset, self).__init__(root, subset='training')\n",
    "        self.len = len(self._walker) // 35\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        index = np.random.randint(0, len(self._walker))\n",
    "        fileid = self._walker[index]\n",
    "        waveform, sample_rate, _, speaker_id, utterance_number = load_speechcommands_item(fileid, self._path)\n",
    "        return waveform, sample_rate, \"unknown\", speaker_id, utterance_number\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The PyTorch Lightning Data Module for KWS\n",
    "\n",
    "`KWSDataModule` cleanly separates the data handling from the model. The data module handles the datasets and dataloaders.\n",
    "\n",
    "We use `torchaudio` `SPEECHCOMMANDS` dataset to load the training, testing and validation sets. \n",
    "\n",
    "A custom `collate_fn` is used to handle the different lengths of the audio samples. The function also converts the wav files into mel spectrograms for the ResNet18 model input layer. A mel spectrogram is a log-mel spectrogram. It is an image that shows the power spectrum of the audio signal in dB. Basically, we convert audio to image. Then, we can use an image classifier like ResNet18."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KWSDataModule(LightningDataModule):\n",
    "    def __init__(self, path, batch_size=128, num_workers=0, n_fft=512, \n",
    "                 n_mels=128, win_length=None, hop_length=256, class_dict={}, \n",
    "                 **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.path = path\n",
    "        self.batch_size = batch_size\n",
    "        self.num_workers = num_workers\n",
    "        self.n_fft = n_fft\n",
    "        self.n_mels = n_mels\n",
    "        self.win_length = win_length\n",
    "        self.hop_length = hop_length\n",
    "        self.class_dict = class_dict\n",
    "\n",
    "    def prepare_data(self):\n",
    "        self.train_dataset = torchaudio.datasets.SPEECHCOMMANDS(self.path,\n",
    "                                                                download=True,\n",
    "                                                                subset='training')\n",
    "\n",
    "        silence_dataset = SilenceDataset(self.path)\n",
    "        unknown_dataset = UnknownDataset(self.path)\n",
    "        self.train_dataset = torch.utils.data.ConcatDataset([self.train_dataset, silence_dataset, unknown_dataset])\n",
    "                                                                \n",
    "        self.val_dataset = torchaudio.datasets.SPEECHCOMMANDS(self.path,\n",
    "                                                              download=True,\n",
    "                                                              subset='validation')\n",
    "        self.test_dataset = torchaudio.datasets.SPEECHCOMMANDS(self.path,\n",
    "                                                               download=True,\n",
    "                                                               subset='testing')                                                    \n",
    "        _, sample_rate, _, _, _ = self.train_dataset[0]\n",
    "        self.sample_rate = sample_rate\n",
    "        self.transform = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate,\n",
    "                                                              n_fft=self.n_fft,\n",
    "                                                              win_length=self.win_length,\n",
    "                                                              hop_length=self.hop_length,\n",
    "                                                              n_mels=self.n_mels,\n",
    "                                                              power=2.0)\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        self.prepare_data()\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.train_dataset,\n",
    "            batch_size=self.batch_size,\n",
    "            num_workers=self.num_workers,\n",
    "            shuffle=True,\n",
    "            pin_memory=True,\n",
    "            collate_fn=self.collate_fn\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.val_dataset,\n",
    "            batch_size=self.batch_size,\n",
    "            num_workers=self.num_workers,\n",
    "            shuffle=True,\n",
    "            pin_memory=True,\n",
    "            collate_fn=self.collate_fn\n",
    "        )\n",
    "    \n",
    "    def test_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.test_dataset,\n",
    "            batch_size=self.batch_size,\n",
    "            num_workers=self.num_workers,\n",
    "            shuffle=True,\n",
    "            pin_memory=True,\n",
    "            collate_fn=self.collate_fn\n",
    "        )\n",
    "\n",
    "    def collate_fn(self, batch):\n",
    "        mels = []\n",
    "        labels = []\n",
    "        wavs = []\n",
    "        for sample in batch:\n",
    "            waveform, sample_rate, label, speaker_id, utterance_number = sample\n",
    "            # ensure that all waveforms are 1sec in length; if not pad with zeros\n",
    "            if waveform.shape[-1] < sample_rate:\n",
    "                waveform = torch.cat([waveform, torch.zeros((1, sample_rate - waveform.shape[-1]))], dim=-1)\n",
    "            elif waveform.shape[-1] > sample_rate:\n",
    "                waveform = waveform[:,:sample_rate]\n",
    "\n",
    "            # mel from power to db\n",
    "            mels.append(ToTensor()(librosa.power_to_db(self.transform(waveform).squeeze().numpy(), ref=np.max)))\n",
    "            labels.append(torch.tensor(self.class_dict[label]))\n",
    "            wavs.append(waveform)\n",
    "\n",
    "        mels = torch.stack(mels)\n",
    "        labels = torch.stack(labels)\n",
    "        wavs = torch.stack(wavs)\n",
    "   \n",
    "        return mels, labels, wavs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The PL LightningModule for KWS\n",
    "\n",
    "The `KWSModel` is a ResNet18 model with first and last layers modified to suport single channel mel spectrogram inputs and 37-category outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KWSModel(LightningModule):\n",
    "    def __init__(self, num_classes=37, epochs=30, lr=0.001, **kwargs):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "        self.model = torchvision.models.resnet18(num_classes=num_classes)\n",
    "        self.model.conv1 = torch.nn.Conv2d(\n",
    "            1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        mels, labels, _ = batch\n",
    "        preds = self.model(mels)\n",
    "        loss = self.hparams.criterion(preds, labels)\n",
    "        return {'loss': loss}\n",
    "\n",
    "    # calls to self.log() are recorded in wandb\n",
    "    def training_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
    "        self.log(\"train_loss\", avg_loss, on_epoch=True)\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        return self.test_step(batch, batch_idx)\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        return self.test_epoch_end(outputs)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        mels, labels, wavs = batch\n",
    "        preds = self.model(mels)\n",
    "        loss = self.hparams.criterion(preds, labels)\n",
    "        acc = accuracy(preds, labels) * 100.\n",
    "        return {\"preds\": preds, 'test_loss': loss, 'test_acc': acc}\n",
    "\n",
    "    def test_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()\n",
    "        avg_acc = torch.stack([x['test_acc'] for x in outputs]).mean()\n",
    "        self.log(\"test_loss\", avg_loss, on_epoch=True, prog_bar=True)\n",
    "        self.log(\"test_acc\", avg_acc, on_epoch=True, prog_bar=True)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            self.parameters(), lr=self.hparams.lr)\n",
    "        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "            optimizer=optimizer, T_max=self.hparams.epochs)\n",
    "        return [optimizer], [lr_scheduler]\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        self.hparams.criterion = torch.nn.CrossEntropyLoss()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Lightning Callback\n",
    "\n",
    "We can instantiate a callback object to perform certain tasks during training. In this case, we log sample audio files, ground truth labels, and predicted labels.\n",
    "\n",
    "We can also `ModelCheckpoint` callback to save the model after each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WandbCallback(Callback):\n",
    "\n",
    "    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):\n",
    "        # log 10 sample audio predictions from the first batch\n",
    "        if batch_idx == 0:\n",
    "            n = 10\n",
    "            mels, labels, wavs = batch\n",
    "            preds = outputs[\"preds\"]\n",
    "            preds = torch.argmax(preds, dim=1)\n",
    "\n",
    "            labels = labels.cpu().numpy()\n",
    "            preds = preds.cpu().numpy()\n",
    "            \n",
    "            wavs = torch.squeeze(wavs, dim=1)\n",
    "            wavs = [ (wav.cpu().numpy()*32768.0).astype(\"int16\") for wav in wavs]\n",
    "            \n",
    "            sample_rate = pl_module.hparams.sample_rate\n",
    "            idx_to_class = pl_module.hparams.idx_to_class\n",
    "            \n",
    "            # log audio samples and predictions as a W&B Table\n",
    "            columns = ['audio', 'mel', 'ground truth', 'prediction']\n",
    "            data = [[wandb.Audio(wav, sample_rate=sample_rate), wandb.Image(mel), idx_to_class[label], idx_to_class[pred]] for wav, mel, label, pred in list(\n",
    "                zip(wavs[:n], mels[:n], labels[:n], preds[:n]))]\n",
    "            wandb_logger.log_table(\n",
    "                key='ResNet18 on KWS using PyTorch Lightning',\n",
    "                columns=columns,\n",
    "                data=data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Program Arguments\n",
    "\n",
    "The default configuration is shown below. \n",
    "\n",
    "We also define `plot_waveform` for plotting the waveform of the audio samples.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_args():\n",
    "    parser = argparse.ArgumentParser()\n",
    "    # model training hyperparameters\n",
    "    parser.add_argument('--batch-size', type=int, default=64, metavar='N',\n",
    "                        help='input batch size for training (default: 64)')\n",
    "    parser.add_argument('--max-epochs', type=int, default=30, metavar='N',\n",
    "                        help='number of epochs to train (default: 30)')\n",
    "    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',\n",
    "                        help='learning rate (default: 0.001)')\n",
    "\n",
    "    # where dataset will be stored\n",
    "    parser.add_argument(\"--path\", type=str, default=\"data/speech_commands/\")\n",
    "\n",
    "    # 35 keywords + silence + unknown\n",
    "    parser.add_argument(\"--num-classes\", type=int, default=37)\n",
    "   \n",
    "    # mel spectrogram parameters\n",
    "    parser.add_argument(\"--n-fft\", type=int, default=1024)\n",
    "    parser.add_argument(\"--n-mels\", type=int, default=128)\n",
    "    parser.add_argument(\"--win-length\", type=int, default=None)\n",
    "    parser.add_argument(\"--hop-length\", type=int, default=512)\n",
    "\n",
    "    # 16-bit fp model to reduce the size\n",
    "    parser.add_argument(\"--precision\", default=16)\n",
    "    parser.add_argument(\"--accelerator\", default='gpu')\n",
    "    parser.add_argument(\"--devices\", default=1)\n",
    "    parser.add_argument(\"--num-workers\", type=int, default=48)\n",
    "\n",
    "    parser.add_argument(\"--no-wandb\", default=False, action='store_true')\n",
    "\n",
    "    args = parser.parse_args(\"\")\n",
    "    return args\n",
    "\n",
    "def plot_waveform(waveform, sample_rate, title=\"Waveform\", xlim=None, ylim=None):\n",
    "    waveform = waveform.numpy()\n",
    "\n",
    "    num_channels, num_frames = waveform.shape\n",
    "    time_axis = torch.arange(0, num_frames) / sample_rate\n",
    "\n",
    "    figure, axes = plt.subplots(num_channels, 1)\n",
    "    if num_channels == 1:\n",
    "        axes = [axes]\n",
    "    for c in range(num_channels):\n",
    "        axes[c].plot(time_axis, waveform[c], linewidth=1)\n",
    "        axes[c].grid(True)\n",
    "        if num_channels > 1:\n",
    "            axes[c].set_ylabel(f'Channel {c+1}')\n",
    "        if xlim:\n",
    "            axes[c].set_xlim(xlim)\n",
    "        if ylim:\n",
    "            axes[c].set_ylim(ylim)\n",
    "    figure.suptitle(title)\n",
    "    plt.show(block=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### KWS Training and Evaluation using `Trainer`\n",
    "\n",
    "The actual training and validation using `Trainer`.\n",
    "\n",
    "Accuracy is about `94.5%` (`94.0%` 16-bit) on the test set. Benchmark accuracy is `88.2%` on the original paper. The state of the art (as of 1 Apr 2022) is `97.0%`. For updated benchmarks, see [paperswithcode.com](https://paperswithcode.com/sota/keyword-spotting-on-google-speech-commands)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KWSModel(\n",
      "  (model): ResNet(\n",
      "    (conv1): Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (relu): ReLU(inplace=True)\n",
      "    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "    (layer1): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer2): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer3): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer4): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "    (fc): Linear(in_features=512, out_features=37, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mrowel\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.4"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>./wandb/run-20221011_092151-90zols1x</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/rowel/pl-kws/runs/90zols1x\" target=\"_blank\">curious-field-10</a></strong> to <a href=\"https://wandb.ai/rowel/pl-kws\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit native Automatic Mixed Precision (AMP)\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:616: UserWarning: Checkpoint directory /home/rowel/github/roatienza/Deep-Learning-Experiments/versions/2022/supervised/python/data/speech_commands/checkpoints exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [1]\n",
      "\n",
      "  | Name  | Type   | Params\n",
      "---------------------------------\n",
      "0 | model | ResNet | 11.2 M\n",
      "---------------------------------\n",
      "11.2 M    Trainable params\n",
      "0         Non-trainable params\n",
      "11.2 M    Total params\n",
      "22.378    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.021261215209960938,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Sanity Checking",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdd7177712d94303a1c5aea9663c1330",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:489: PossibleUserWarning: Your `val_dataloader`'s sampler has shuffling enabled, it is strongly recommended that you turn shuffling off for val/test/predict dataloaders.\n",
      "  rank_zero_warn(\n",
      "Exception in thread StatsThr:\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/threading.py\", line 973, in _bootstrap_inner\n",
      "    self.run()\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/threading.py\", line 910, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/sdk/internal/stats.py\", line 137, in _thread_body\n",
      "    stats = self.stats()\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/sdk/internal/stats.py\", line 183, in stats\n",
      "    handle = pynvml.nvmlDeviceGetHandleByIndex(i)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/vendor/pynvml/pynvml.py\", line 2232, in nvmlDeviceGetHandleByIndex\n",
      "    _nvmlCheckReturn(ret)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/vendor/pynvml/pynvml.py\", line 1015, in _nvmlCheckReturn\n",
      "    raise NVMLError(ret)\n",
      "wandb.vendor.pynvml.pynvml.NVMLError_GpuIsLost: GPU is lost\n"
     ]
    }
   ],
   "source": [
    "\n",
    "if __name__ == \"__main__\":\n",
    "\n",
    "    args = get_args()\n",
    "    CLASSES = ['silence', 'unknown', 'backward', 'bed', 'bird', 'cat', 'dog', 'down', 'eight', 'five', 'follow',\n",
    "               'forward', 'four', 'go', 'happy', 'house', 'learn', 'left', 'marvin', 'nine', 'no',\n",
    "               'off', 'on', 'one', 'right', 'seven', 'sheila', 'six', 'stop', 'three',\n",
    "               'tree', 'two', 'up', 'visual', 'wow', 'yes', 'zero']\n",
    "    \n",
    "    # make a dictionary from CLASSES to integers\n",
    "    CLASS_TO_IDX = {c: i for i, c in enumerate(CLASSES)}\n",
    "\n",
    "    if not os.path.exists(args.path):\n",
    "        os.makedirs(args.path, exist_ok=True)\n",
    "\n",
    "    model = KWSModel(num_classes=args.num_classes, epochs=args.max_epochs, lr=args.lr)\n",
    "    print(model)\n",
    "    datamodule = KWSDataModule(batch_size=args.batch_size, num_workers=args.num_workers,\n",
    "                               path=args.path, n_fft=args.n_fft, n_mels=args.n_mels,\n",
    "                               win_length=args.win_length, hop_length=args.hop_length,\n",
    "                               class_dict=CLASS_TO_IDX)\n",
    "    datamodule.setup()\n",
    "\n",
    "    # wandb is a great way to debug and visualize this model\n",
    "    wandb_logger = WandbLogger(project=\"pl-kws\")\n",
    "\n",
    "    model_checkpoint = ModelCheckpoint(\n",
    "        dirpath=os.path.join(args.path, \"checkpoints\"),\n",
    "        filename=\"resnet18-kws-best-acc\",\n",
    "        save_top_k=1,\n",
    "        verbose=True,\n",
    "        monitor='test_acc',\n",
    "        mode='max',\n",
    "    )\n",
    "    idx_to_class = {v: k for k, v in CLASS_TO_IDX.items()}\n",
    "    trainer = Trainer(accelerator=args.accelerator,\n",
    "                      devices=args.devices,\n",
    "                      precision=args.precision,\n",
    "                      max_epochs=args.max_epochs,\n",
    "                      logger=wandb_logger if not args.no_wandb else None,\n",
    "                      callbacks=[model_checkpoint, WandbCallback() if not args.no_wandb else None])\n",
    "    model.hparams.sample_rate = datamodule.sample_rate\n",
    "    model.hparams.idx_to_class = idx_to_class\n",
    "    trainer.fit(model, datamodule=datamodule)\n",
    "    trainer.test(model, datamodule=datamodule)\n",
    "\n",
    "    wandb.finish()\n",
    "    #trainer.save_checkpoint('../mnist/checkpoint.ckpt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground Truth: sheila, Prediction: sheila\n"
     ]
    }
   ],
   "source": [
    "# https://pytorch-lightning.readthedocs.io/en/stable/common/production_inference.html\n",
    "model = model.load_from_checkpoint(os.path.join(\n",
    "    args.path, \"checkpoints\", \"resnet18-kws-best-acc.ckpt\"))\n",
    "model.eval()\n",
    "script = model.to_torchscript()\n",
    "\n",
    "# save for use in production environment\n",
    "model_path = os.path.join(args.path, \"checkpoints\",\n",
    "                          \"resnet18-kws-best-acc.pt\")\n",
    "torch.jit.save(script, model_path)\n",
    "\n",
    "# list wav files given a folder\n",
    "label = CLASSES[2:]\n",
    "label = np.random.choice(label)\n",
    "path = os.path.join(args.path, \"SpeechCommands/speech_commands_v0.02/\")\n",
    "path = os.path.join(path, label)\n",
    "wav_files = [os.path.join(path, f)\n",
    "             for f in os.listdir(path) if f.endswith('.wav')]\n",
    "# select random wav file\n",
    "wav_file = np.random.choice(wav_files)\n",
    "waveform, sample_rate = torchaudio.load(wav_file)\n",
    "transform = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate,\n",
    "                                                 n_fft=args.n_fft,\n",
    "                                                 win_length=args.win_length,\n",
    "                                                 hop_length=args.hop_length,\n",
    "                                                 n_mels=args.n_mels,\n",
    "                                                 power=2.0)\n",
    "\n",
    "mel = ToTensor()(librosa.power_to_db(\n",
    "    transform(waveform).squeeze().numpy(), ref=np.max))\n",
    "mel = mel.unsqueeze(0)\n",
    "scripted_module = torch.jit.load(model_path)\n",
    "pred = torch.argmax(scripted_module(mel), dim=1)\n",
    "print(f\"Ground Truth: {label}, Prediction: {idx_to_class[pred.item()]}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('voice')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "975196f818cef4d8418ca9caaf131be767e393cdee268c01707e7bb8d73c6de2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
